{
  "cells": [
    {
      "metadata": {
        "id": "zWgEkOAO9OVz"
      },
      "cell_type": "markdown",
      "source": [
        "# An example implementing the three key questions\n",
        "\n",
        "March and Olsen (2011) posit that humans generally act as though they choose their actions by answering three key questions:\n",
        "\n",
        "1. What kind of situation is this?\n",
        "2. What kind of person am I?\n",
        "3. What does a person such as I do in a situation such as this?\n",
        "\n",
        "The agents used in this example implement these components. The premise of the simulation is that 4 friends are stuck at a snowed in pub. Also, it has just come to light that Alice stole Bob's car and crashed it.\n"
      ]
    },
    {
      "metadata": {
        "id": "esFO3miE3s41"
      },
      "cell_type": "markdown",
      "source": [
        "\u003ca href=\"https://colab.research.google.com/github/google-deepmind/concordia/blob/main/examples/deprecated/example_with_scenes.ipynb\" target=\"_parent\"\u003e\u003cimg src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/\u003e\u003c/a\u003e"
      ]
    },
    {
      "metadata": {
        "id": "J2TwJrZ08wXz"
      },
      "cell_type": "markdown",
      "source": [
        "## Setup and imports"
      ]
    },
    {
      "metadata": {
        "id": "tkCJgbFFvaAB"
      },
      "cell_type": "code",
      "source": [
        "# @title Colab-specific setup (use a CodeSpace to avoid the need for this).\n",
        "try:\n",
        "  %env COLAB_RELEASE_TAG\n",
        "except:\n",
        "  pass  # Not running in colab.\n",
        "else:\n",
        "  %pip install --ignore-requires-python --requirement 'https://raw.githubusercontent.com/google-deepmind/concordia/main/examples/requirements.in' 'git+https://github.com/google-deepmind/concordia.git#egg=gdm-concordia'\n",
        "  %pip list"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "metadata": {
        "id": "-qLG5ExLqpWa"
      },
      "cell_type": "code",
      "source": [
        "# @title Imports\n",
        "\n",
        "import collections\n",
        "from collections.abc import Callable, Mapping, Sequence\n",
        "import concurrent.futures\n",
        "import datetime\n",
        "\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "import sentence_transformers\n",
        "\n",
        "from IPython import display\n",
        "\n",
        "from concordia.agents.deprecated import deprecated_agent as basic_agent\n",
        "from concordia.components import deprecated as generic_components\n",
        "from concordia.components.agent.deprecated import to_be_deprecated as agent_components\n",
        "from concordia.associative_memory.deprecated import associative_memory\n",
        "from concordia.associative_memory.deprecated import blank_memories\n",
        "from concordia.associative_memory.deprecated import formative_memories\n",
        "from concordia.associative_memory.deprecated import importance_function\n",
        "from concordia.clocks import game_clock\n",
        "from concordia.components.game_master import deprecated as gm_components\n",
        "from concordia.environment.deprecated import game_master\n",
        "from concordia.environment.deprecated.scenes import runner\n",
        "from concordia.contrib.language_models.openai import gpt_model\n",
        "from concordia.language_model import language_model\n",
        "from concordia.deprecated.metrics import goal_achievement\n",
        "from concordia.deprecated.metrics import common_sense_morality\n",
        "from concordia.deprecated.metrics import opinion_of_others\n",
        "from concordia.thought_chains.deprecated import thought_chains as thought_chains_lib\n",
        "from concordia.typing.deprecated import agent as agent_lib\n",
        "from concordia.typing.deprecated import component\n",
        "from concordia.typing.deprecated import scene as scene_lib\n",
        "from concordia.utils import html as html_lib\n",
        "from concordia.utils.deprecated import measurements as measurements_lib\n",
        "from concordia.utils import plotting\n",
        "\n",
        "import termcolor"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "metadata": {
        "id": "I3OtW8flCJSC"
      },
      "cell_type": "code",
      "source": [
        "# Setup sentence encoder\n",
        "st_model = sentence_transformers.SentenceTransformer(\n",
        "    'sentence-transformers/all-mpnet-base-v2')\n",
        "embedder = lambda x: st_model.encode(x, show_progress_bar=False)"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "metadata": {
        "id": "cugwvFIKv5AS"
      },
      "cell_type": "code",
      "source": [
        "# @title Language Model - pick your model and provide keys\n",
        "\n",
        "# By default this colab uses GPT-4, so you must provide an API key.\n",
        "# Note that it is also possible to use local models or other API models,\n",
        "# simply replace this cell with the correct initialization for the model\n",
        "# you want to use.\n",
        "GPT_API_KEY = '' #@param {type: 'string'}\n",
        "GPT_MODEL_NAME = 'gpt-4o' #@param {type: 'string'}\n",
        "\n",
        "if not GPT_API_KEY:\n",
        "  raise ValueError('GPT_API_KEY is required.')\n",
        "\n",
        "model = gpt_model.GptLanguageModel(api_key=GPT_API_KEY,\n",
        "                                   model_name=GPT_MODEL_NAME)\n"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "metadata": {
        "id": "13y2g8Vq61Vw"
      },
      "cell_type": "markdown",
      "source": [
        "## Configuring the generic knowledge shared by players and GM."
      ]
    },
    {
      "metadata": {
        "id": "Q89w-OrDf7c_"
      },
      "cell_type": "code",
      "source": [
        "# @title Shared memories are memories that all players and GM share.\n",
        "\n",
        "DECISION_SCENE_TYPE = 'decision'\n",
        "\n",
        "shared_memories = [\n",
        "    'There is a pub called The Sundrop Saloon.',\n",
        "    'It is impossible to leave the Sundrop Saloon, since it is snowed in.',\n",
        "    \"Alice stole Bob's car and crashed it.\",\n",
        "    ('Alice, Bob, and Dorothy always spend their evenings at the ' +\n",
        "     'Sundrop Saloon.'),\n",
        "]\n",
        "\n",
        "# The shared context will be used for the NPC context. It reflects general\n",
        "# knowledge and is possessed by all characters.\n",
        "shared_context = model.sample_text(\n",
        "    'Summarize the following passage in a concise and insightful fashion:\\n'\n",
        "    + '\\n'.join(shared_memories)\n",
        "    + '\\n'\n",
        "    + 'Summary:'\n",
        ")\n",
        "print(shared_context)"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "metadata": {
        "id": "rcEWUWFqrYVO"
      },
      "cell_type": "code",
      "source": [
        "PERSONALITY_TRAITS = [\n",
        "    \"Aggressive\",\n",
        "    \"Optimistic\",\n",
        "    \"Kind\",\n",
        "    \"Resilient\",\n",
        "    \"Humorous\",\n",
        "    \"Empathetic\",\n",
        "    \"Ambitious\",\n",
        "    \"Honest\",\n",
        "    \"Loyal\",\n",
        "    \"Pessimistic\",\n",
        "    \"Arrogant\",\n",
        "    \"Impulsive\",\n",
        "    \"Jealous\",\n",
        "    \"Manipulative\",\n",
        "    \"Creative\",\n",
        "    \"Analytical\",\n",
        "    \"Confident\",\n",
        "    \"Passionate\",\n",
        "    \"Anxious\",\n",
        "    \"Closed-minded\",\n",
        "    \"Deceitful\",\n",
        "    \"Insecure\",\n",
        "    \"Irresponsible\",\n",
        "    \"Vindictive\",\n",
        "    \"Curious\",\n",
        "    \"Energetic\",\n",
        "    \"Sarcastic\",\n",
        "]\n",
        "\n",
        "def get_trait() -\u003e str:\n",
        "  return np.random.choice(PERSONALITY_TRAITS)"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "metadata": {
        "id": "c3jmyBMFf88V"
      },
      "cell_type": "code",
      "source": [
        "#@title Creating character backgrounds, goals and traits. Modify to explore how it influences the outcomes\n",
        "scenario_premise = [\n",
        "    'Alice stole Bob\\'s car and crashed it.',\n",
        "    'There is a blizzard',\n",
        "]\n",
        "player_configs = [\n",
        "    formative_memories.AgentConfig(\n",
        "        name='Alice',\n",
        "        gender='female',\n",
        "        goal=('Alice wants to live her life without bothing to think about ' +\n",
        "              'others.'),\n",
        "        context=shared_context,\n",
        "        traits=f'traits: Aggressive, {get_trait()}, and {get_trait()}',\n",
        "        extras={\n",
        "            'player_specific_memories': [\n",
        "                f'Alice thinks Bob is {get_trait()}.',\n",
        "                'Alice is a great driver, much better than Bob.',\n",
        "                'Alice accidentally crashed Bob\\'s car, but she\\'s not sorry.',\n",
        "                'Alice is very impatient and dislikes long conversations.',\n",
        "            ]\n",
        "        },\n",
        "    ),\n",
        "    formative_memories.AgentConfig(\n",
        "        name='Bob',\n",
        "        gender='male',\n",
        "        goal='Bob wants Alice to pay for his car.',\n",
        "        context=shared_context,\n",
        "        traits=f'traits: {get_trait()}, {get_trait()}, and {get_trait()}',\n",
        "        extras={\n",
        "            'player_specific_memories': [\n",
        "                'Bob had to save for years to afford his beloved car.'\n",
        "            ]\n",
        "        },\n",
        "    ),\n",
        "    formative_memories.AgentConfig(\n",
        "        name='Dorothy',\n",
        "        gender='female',\n",
        "        goal=(\n",
        "            'Dorothy wants to create a conflict between Bob and Alice for fun.'\n",
        "        ),\n",
        "        context=shared_context,\n",
        "        traits=f'traits: {get_trait()}, {get_trait()}, and {get_trait()}',\n",
        "        extras={\n",
        "            'player_specific_memories': [\n",
        "                'Dorothy grew up in Riverbend.',\n",
        "                'Dorothy hates boring people.',\n",
        "            ]\n",
        "        },\n",
        "    ),\n",
        "]\n",
        "num_players = len(player_configs)\n",
        "\n",
        "player_configs_dict = {player.name: player for player in player_configs}"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "metadata": {
        "id": "J46YgxCKf7aZ"
      },
      "cell_type": "code",
      "source": [
        "#@title Make the clock\n",
        "time_step = datetime.timedelta(minutes=10)\n",
        "SETUP_TIME = datetime.datetime(hour=20, year=2024, month=10, day=1)\n",
        "\n",
        "START_TIME = datetime.datetime(hour=16, year=2024, month=10, day=2)\n",
        "PUB_TIME = datetime.datetime(hour=23, year=2024, month=10, day=2)\n",
        "clock = game_clock.MultiIntervalClock(\n",
        "    start=SETUP_TIME,\n",
        "    step_sizes=[time_step, datetime.timedelta(seconds=10)])"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "metadata": {
        "id": "h4_gUs6wrjPM"
      },
      "cell_type": "code",
      "source": [
        "#@title Importance models\n",
        "importance_model = importance_function.AgentImportanceModel(model)\n",
        "importance_model_gm = importance_function.ConstantImportanceModel()"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "metadata": {
        "id": "qt8CK2mMbD7q"
      },
      "cell_type": "markdown",
      "source": [
        "## Configure and build the players\n",
        "\n",
        "---\n",
        "\n"
      ]
    },
    {
      "metadata": {
        "id": "CrmDfTNHCVXC"
      },
      "cell_type": "code",
      "source": [
        "blank_memory_factory = blank_memories.MemoryFactory(\n",
        "    model=model,\n",
        "    embedder=embedder,\n",
        "    importance=importance_model.importance,\n",
        "    clock_now=clock.now,\n",
        ")\n",
        "\n",
        "formative_memory_factory = formative_memories.FormativeMemoryFactory(\n",
        "    model=model,\n",
        "    shared_memories=shared_memories,\n",
        "    blank_memory_factory_call=blank_memory_factory.make_blank_memory,\n",
        ")"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "metadata": {
        "id": "4aS2sY22B1JQ"
      },
      "cell_type": "code",
      "source": [
        "#@title Agent architecture definition\n",
        "\n",
        "def build_agent(agent_config,\n",
        "                player_names: list[str],\n",
        "                measurements: measurements_lib.Measurements | None = None):\n",
        "\n",
        "  mem = formative_memory_factory.make_memories(agent_config)\n",
        "  agent_name = agent_config.name\n",
        "\n",
        "  instructions = generic_components.constant.ConstantComponent(\n",
        "      state=(\n",
        "          f'The instructions for how to play the role of {agent_name} are as '\n",
        "          'follows. This is a social science experiment studying how well you '\n",
        "          f'play the role of a character named {agent_name}. The experiment '\n",
        "          'is structured as a tabletop roleplaying game (like dungeons and '\n",
        "          'dragons). However, in this case it is a serious social science '\n",
        "          'experiment and simulation. The goal is to be realistic. It is '\n",
        "          f'important to play the role of a person like {agent_name} as '\n",
        "          f'accurately as possible, i.e., by responding in ways that you think '\n",
        "          f'it is likely a person like {agent_name} would respond, and taking '\n",
        "          f'into account all information about {agent_name} that you have. '\n",
        "          'Always use third-person limited perspective.'\n",
        "      ),\n",
        "      name='role playing instructions\\n')\n",
        "\n",
        "  time = generic_components.report_function.ReportFunction(\n",
        "      name='Current time',\n",
        "      function=clock.current_time_interval_str,\n",
        "  )\n",
        "\n",
        "  current_obs = agent_components.observation.Observation(\n",
        "      agent_name=agent_config.name,\n",
        "      clock_now=clock.now,\n",
        "      memory=mem,\n",
        "      timeframe=clock.get_step_size(),\n",
        "      component_name='current observations',\n",
        "  )\n",
        "  somatic_state = agent_components.somatic_state.SomaticState(\n",
        "      model=model,\n",
        "      memory=mem,\n",
        "      agent_name=agent_config.name,\n",
        "      clock_now=clock.now,\n",
        "  )\n",
        "  summary_obs = agent_components.observation.ObservationSummary(\n",
        "      agent_name=agent_config.name,\n",
        "      model=model,\n",
        "      clock_now=clock.now,\n",
        "      memory=mem,\n",
        "      components=[current_obs, somatic_state],\n",
        "      timeframe_delta_from=datetime.timedelta(hours=4),\n",
        "      timeframe_delta_until=datetime.timedelta(hours=1),\n",
        "      component_name='summary of observations',\n",
        "  )\n",
        "\n",
        "  self_perception = agent_components.self_perception.SelfPerception(\n",
        "      name=f'answer to what kind of person is {agent_config.name}',\n",
        "      model=model,\n",
        "      memory=mem,\n",
        "      agent_name=agent_config.name,\n",
        "      clock_now=clock.now,\n",
        "  )\n",
        "  relevant_memories = agent_components.all_similar_memories.AllSimilarMemories(\n",
        "      name='relevant memories',\n",
        "      model=model,\n",
        "      memory=mem,\n",
        "      agent_name=agent_name,\n",
        "      components=[summary_obs, self_perception],\n",
        "      clock_now=clock.now,\n",
        "      num_memories_to_retrieve=10,\n",
        "  )\n",
        "  situation_perception = agent_components.situation_perception.SituationPerception(\n",
        "      name=(f'answer to what kind of situation is {agent_config.name} in ' +\n",
        "            'right now'),\n",
        "      model=model,\n",
        "      memory=mem,\n",
        "      agent_name=agent_config.name,\n",
        "      components=[current_obs, somatic_state, summary_obs],\n",
        "      clock_now=clock.now,\n",
        "  )\n",
        "  person_by_situation = agent_components.person_by_situation.PersonBySituation(\n",
        "      name=(f'answer to what would a person like {agent_config.name} do in a ' +\n",
        "            'situation like this'),\n",
        "      model=model,\n",
        "      memory=mem,\n",
        "      agent_name=agent_config.name,\n",
        "      clock_now=clock.now,\n",
        "      components=[self_perception, situation_perception],\n",
        "      verbose=True,\n",
        "  )\n",
        "\n",
        "  persona = generic_components.sequential.Sequential(\n",
        "      name='persona',\n",
        "      components=[\n",
        "          self_perception,\n",
        "          situation_perception,\n",
        "          person_by_situation,\n",
        "      ]\n",
        "  )\n",
        "\n",
        "  reflection = agent_components.dialectical_reflection.DialecticalReflection(\n",
        "        name='reflection',\n",
        "        model=model,\n",
        "        memory=mem,\n",
        "        agent_name=agent_config.name,\n",
        "        intuition_components=[self_perception],\n",
        "        thinking_components=[persona],\n",
        "        clock_now=clock.now,\n",
        "        num_memories_to_retrieve=3,\n",
        "        verbose=True,\n",
        "    )\n",
        "\n",
        "  initial_goal_component = generic_components.constant.ConstantComponent(\n",
        "      state=agent_config.goal, name='overarching goal')\n",
        "  plan = agent_components.plan.SimPlan(\n",
        "      model,\n",
        "      mem,\n",
        "      agent_config.name,\n",
        "      clock_now=clock.now,\n",
        "      components=[instructions, initial_goal_component, somatic_state, persona],\n",
        "      goal=person_by_situation,\n",
        "      horizon='the next hour',\n",
        "      verbose=True,\n",
        "  )\n",
        "\n",
        "  goal_metric = goal_achievement.GoalAchievementMetric(\n",
        "      model=model,\n",
        "      player_name=agent_config.name,\n",
        "      player_goal=agent_config.goal,\n",
        "      clock=clock,\n",
        "      name='Goal Achievement',\n",
        "      measurements=measurements,\n",
        "      channel='goal_achievement',\n",
        "      verbose=False,\n",
        "  )\n",
        "  morality_metric = common_sense_morality.CommonSenseMoralityMetric(\n",
        "      model=model,\n",
        "      player_name=agent_config.name,\n",
        "      clock=clock,\n",
        "      name='Morality',\n",
        "      verbose=False,\n",
        "      measurements=measurements,\n",
        "      channel='common_sense_morality',\n",
        "  )\n",
        "  agent = basic_agent.BasicAgent(\n",
        "      model,\n",
        "      agent_name=agent_config.name,\n",
        "      clock=clock,\n",
        "      verbose=False,\n",
        "      components=[instructions,\n",
        "                  persona,\n",
        "                  plan,\n",
        "                  relevant_memories,\n",
        "                  reflection,\n",
        "                  time,\n",
        "                  current_obs,\n",
        "                  goal_metric,\n",
        "                  morality_metric],\n",
        "      update_interval = time_step\n",
        "  )\n",
        "  reputation_metric = opinion_of_others.OpinionOfOthersMetric(\n",
        "      model=model,\n",
        "      player_name=agent_config.name,\n",
        "      player_names=player_names,\n",
        "      context_fn=agent.state,\n",
        "      clock=clock,\n",
        "      name='Opinion',\n",
        "      verbose=False,\n",
        "      measurements=measurements,\n",
        "      channel='opinion_of_others',\n",
        "      question='What is {opining_player}\\'s opinion of {of_player}?',\n",
        "  )\n",
        "  agent.add_component(reputation_metric)\n",
        "\n",
        "  for extra_memory in agent_config.extras['player_specific_memories']:\n",
        "    agent.observe(extra_memory)\n",
        "  return agent, mem\n"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "metadata": {
        "id": "5RU3ZV4oIknW"
      },
      "cell_type": "code",
      "source": [
        "player_configs = player_configs[:num_players]\n",
        "player_names = [player.name for player in player_configs][:num_players]\n",
        "measurements = measurements_lib.Measurements()\n",
        "\n",
        "players = []\n",
        "memories = {}\n",
        "with concurrent.futures.ThreadPoolExecutor(max_workers=num_players) as pool:\n",
        "  for agent, mem in pool.map(build_agent,\n",
        "                             player_configs[:num_players],\n",
        "                             # All players get the same `player_names`.\n",
        "                             [player_names] * num_players,\n",
        "                             # All players get the same `measurements` object.\n",
        "                             [measurements] * num_players):\n",
        "    players.append(agent)\n",
        "    memories[agent.name] = mem"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "metadata": {
        "id": "2vt8ggYUrW8M"
      },
      "cell_type": "markdown",
      "source": [
        "## Build GM"
      ]
    },
    {
      "metadata": {
        "id": "3W65kHOKQwrv"
      },
      "cell_type": "code",
      "source": [
        "game_master_memory = associative_memory.AssociativeMemory(\n",
        "   sentence_embedder=embedder,\n",
        "   importance=importance_model_gm.importance,\n",
        "   clock=clock.now)"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "metadata": {
        "id": "-cxivChc633z"
      },
      "cell_type": "code",
      "source": [
        "# @title Create components of the Game Master\n",
        "player_names = [player.name for player in players]\n",
        "\n",
        "scenario_knowledge = generic_components.constant.ConstantComponent(\n",
        "    state=' '.join(shared_memories),\n",
        "    name='Background')\n",
        "\n",
        "player_status = gm_components.player_status.PlayerStatus(\n",
        "    clock_now=clock.now,\n",
        "    model=model,\n",
        "    memory=game_master_memory,\n",
        "    player_names=player_names)\n",
        "\n",
        "convo_externality = gm_components.conversation.Conversation(\n",
        "    players=players,\n",
        "    model=model,\n",
        "    memory=game_master_memory,\n",
        "    clock=clock,\n",
        "    burner_memory_factory=blank_memory_factory,\n",
        "    components=[player_status],\n",
        "    allow_self_talk=True,\n",
        "    cap_nonplayer_characters=3,\n",
        "    shared_context=shared_context,\n",
        "    verbose=False,\n",
        ")\n",
        "\n",
        "direct_effect_externality = gm_components.direct_effect.DirectEffect(\n",
        "    players=players,\n",
        "    model=model,\n",
        "    memory=game_master_memory,\n",
        "    clock_now=clock.now,\n",
        "    verbose=False,\n",
        "    components=[player_status]\n",
        ")\n",
        "\n",
        "relevant_events = gm_components.relevant_events.RelevantEvents(\n",
        "    clock.now, model, game_master_memory)\n",
        "time_display = gm_components.time_display.TimeDisplay(clock)\n"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "metadata": {
        "id": "rsSaBbv55qGy"
      },
      "cell_type": "code",
      "source": [
        "# @title Create the game master's thought chain\n",
        "account_for_agency_of_others = thought_chains_lib.AccountForAgencyOfOthers(\n",
        "    model=model, players=players, verbose=False)\n",
        "thought_chain = [\n",
        "    thought_chains_lib.extract_direct_quote,\n",
        "    thought_chains_lib.attempt_to_most_likely_outcome,\n",
        "    thought_chains_lib.result_to_effect_caused_by_active_player,\n",
        "    account_for_agency_of_others,\n",
        "    thought_chains_lib.restore_direct_quote,\n",
        "]"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "metadata": {
        "id": "d_R2BVNOsAwa"
      },
      "cell_type": "code",
      "source": [
        "# @title Create the game master object\n",
        "env = game_master.GameMaster(\n",
        "    model=model,\n",
        "    memory=game_master_memory,\n",
        "    clock=clock,\n",
        "    players=players,\n",
        "    update_thought_chain=thought_chain,\n",
        "    components=[\n",
        "        scenario_knowledge,\n",
        "        player_status,\n",
        "        convo_externality,\n",
        "        direct_effect_externality,\n",
        "        relevant_events,\n",
        "        time_display,\n",
        "    ],\n",
        "    randomise_initiative=True,\n",
        "    player_observes_event=False,\n",
        "    verbose=True,\n",
        ")"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "metadata": {
        "id": "ybP2PXj4jnlO"
      },
      "cell_type": "code",
      "source": [
        "# @title Create the decision scene logic component\n",
        "\n",
        "class Drinking(component.Component):\n",
        "  \"\"\"Define an example component to track a choice about drinking.\n",
        "  \"\"\"\n",
        "\n",
        "  def __init__(\n",
        "      self,\n",
        "      model: language_model.LanguageModel,\n",
        "      memory: associative_memory.AssociativeMemory,\n",
        "      drinking_option: str,\n",
        "      resolution_scene: str,\n",
        "      players: Sequence[basic_agent.BasicAgent],\n",
        "      clock_now: Callable[[], datetime.datetime],\n",
        "      name: str = 'free drinks',\n",
        "      verbose: bool = False,\n",
        "  ):\n",
        "    \"\"\"Initialize an example component to track a choice about drinking.\n",
        "\n",
        "    Args:\n",
        "      model: a language model\n",
        "      memory: an associative memory\n",
        "      drinking_option: which option choice constitutes drinking\n",
        "      resolution_scene: on which scene type should this component be updated\n",
        "        after the event, i.e. when to check the joint action and compute results\n",
        "      players: sequence of agents (a superset of the active players)\n",
        "      clock_now: Function to call to get current time.\n",
        "      name: name of this component e.g. Possessions, Account, Property, etc\n",
        "      verbose: whether to print the full update chain of thought or not\n",
        "    \"\"\"\n",
        "    self._model = model\n",
        "    self._memory = memory\n",
        "    self._drinking_option = drinking_option\n",
        "    self._players = players\n",
        "    self._clock_now = clock_now\n",
        "    self._name = name\n",
        "    self._verbose = verbose\n",
        "\n",
        "    self._history = []\n",
        "    self._state = ''\n",
        "    self._partial_states = {player.name: '' for player in self._players}\n",
        "\n",
        "    self._resolution_scene = resolution_scene\n",
        "    self._current_scene = gm_components.current_scene.CurrentScene(\n",
        "        name='current scene type',\n",
        "        memory=self._memory,\n",
        "        clock_now=self._clock_now,\n",
        "        verbose=self._verbose,\n",
        "    )\n",
        "\n",
        "    self.reset()\n",
        "    self.update()\n",
        "\n",
        "  def reset(self) -\u003e None:\n",
        "    pass\n",
        "\n",
        "  def name(self) -\u003e str:\n",
        "    \"\"\"Returns the name of this component.\"\"\"\n",
        "    return self._name\n",
        "\n",
        "  def get_last_log(self):\n",
        "    if self._history:\n",
        "      return self._history[-1].copy()\n",
        "\n",
        "  def get_history(self):\n",
        "    return self._history.copy()\n",
        "\n",
        "  def state(self) -\u003e str:\n",
        "    return self._state\n",
        "\n",
        "  def partial_state(\n",
        "      self,\n",
        "      player_name: str,\n",
        "  ) -\u003e str:\n",
        "    \"\"\"Return a player-specific view of the component's state.\"\"\"\n",
        "    return self._partial_states[player_name]\n",
        "\n",
        "  def update(self) -\u003e None:\n",
        "    self._current_scene.update()\n",
        "\n",
        "  def _binarize_joint_action(\n",
        "      self,\n",
        "      joint_action: Mapping[str, str]) -\u003e Mapping[str, int]:\n",
        "    binary_joint_action = {name: act == self._drinking_option\n",
        "                           for name, act in joint_action.items()}\n",
        "    return binary_joint_action\n",
        "\n",
        "  def update_before_event(self, player_action_attempt: str) -\u003e None:\n",
        "    # Get the current scene type.\n",
        "    current_scene_type = self._current_scene.state()\n",
        "    if current_scene_type == self._resolution_scene:\n",
        "      # `player_action_attempt` is formatted as \"name: attempt\".\n",
        "      player_name, choice_str = player_action_attempt.split(': ')\n",
        "      if choice_str == self._drinking_option:\n",
        "        self._partial_states[player_name] = (\n",
        "            'Woah! That\\'s one strong drink!')\n",
        "      else:\n",
        "        self._partial_states[player_name] = '...'\n",
        "\n",
        "      if self._verbose:\n",
        "        print(termcolor.colored(self.state(), 'yellow'))\n",
        "\n",
        "    self._state = ' -- '.join(\n",
        "          [f'{name}: {player_state}' for name, player_state\n",
        "           in self._partial_states.items()])\n",
        "\n",
        "  def update_after_event(\n",
        "      self,\n",
        "      unused_event_statement: str,\n",
        "  ) -\u003e None:\n",
        "    update_log = {\n",
        "        'date': self._clock_now(),\n",
        "        'Summary': self.name(),\n",
        "        'Drinks': self.state(),\n",
        "    }\n",
        "    self._memory.extend([self._state,])\n",
        "    self._history.append(update_log)\n"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "metadata": {
        "id": "ZN78yxDXoRFk"
      },
      "cell_type": "code",
      "source": [
        "vodka_choice = Drinking(\n",
        "    model=model,\n",
        "    memory=game_master_memory,\n",
        "    drinking_option='yes',\n",
        "    resolution_scene=DECISION_SCENE_TYPE,\n",
        "    players=players,\n",
        "    clock_now=clock.now,\n",
        "    verbose=True,\n",
        ")"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "metadata": {
        "id": "xoNcPv_gjZgo"
      },
      "cell_type": "code",
      "source": [
        "# @title Create the decision scene game master\n",
        "\n",
        "decision_action_spec = agent_lib.choice_action_spec(\n",
        "    call_to_action=('Would {agent_name} drink it?'),\n",
        "    options=('no', 'yes'),\n",
        "    tag='decision',\n",
        ")\n",
        "decision_env = game_master.GameMaster(\n",
        "    model=model,\n",
        "    memory=game_master_memory,\n",
        "    clock=clock,\n",
        "    name='Decision Environment',\n",
        "    players=players,\n",
        "    components=[vodka_choice],\n",
        "    action_spec=decision_action_spec,\n",
        "    update_thought_chain=[thought_chains_lib.identity],\n",
        "    randomise_initiative=True,\n",
        "    player_observes_event=False,\n",
        "    concurrent_externalities=False,\n",
        "    verbose=True,\n",
        ")"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "metadata": {
        "id": "LXykV_TdwfKq"
      },
      "cell_type": "code",
      "source": [
        "#@title Scenes\n",
        "\n",
        "hour_increment = datetime.timedelta(hours=1)\n",
        "\n",
        "car_crash_scene_premise = (\n",
        "    'Alice just emerged alone from the wreckage of Bob\\'s car. She crashed ' +\n",
        "    'it. It\\'s very lucky that no one was hurt. Bob will be so mad when he ' +\n",
        "    'finds out.. especially since Alice had \"borrowed\" his car without asking..'\n",
        ")\n",
        "pub_scene_premise = ('Alice, Bob, and Dorothy are now at the ' +\n",
        "                    'The Sundrop Saloon, and unfortunately trapped there due ' +\n",
        "                    'to the blizzard. This will be a long night.')\n",
        "decision_scene_premise = ('The bartender says, \"since we are stuck here, ' +\n",
        "                          'would any of you like a shot of vodka? It\\'s on ' +\n",
        "                          'the house!\"')\n",
        "\n",
        "scene_specs = {\n",
        "    'car_crash': scene_lib.SceneTypeSpec(\n",
        "        name='car_crash',\n",
        "        premise={\n",
        "            'Alice': [car_crash_scene_premise],\n",
        "        },\n",
        "    ),\n",
        "    'pub': scene_lib.SceneTypeSpec(\n",
        "        name='pub',\n",
        "        premise={\n",
        "            'Alice': [pub_scene_premise],\n",
        "            'Bob': [pub_scene_premise],\n",
        "            'Dorothy': [pub_scene_premise],\n",
        "        },\n",
        "    ),\n",
        "    DECISION_SCENE_TYPE: scene_lib.SceneTypeSpec(\n",
        "        name=DECISION_SCENE_TYPE,\n",
        "        premise={\n",
        "            'Alice': [decision_scene_premise],\n",
        "            'Bob': [decision_scene_premise],\n",
        "            'Dorothy': [decision_scene_premise],\n",
        "        },\n",
        "        action_spec=decision_action_spec,\n",
        "        override_game_master=decision_env,\n",
        "    ),\n",
        "}\n",
        "\n",
        "scenes = [\n",
        "    scene_lib.SceneSpec(\n",
        "        scene_type=scene_specs['car_crash'],\n",
        "        start_time=START_TIME,\n",
        "        participant_configs=[player_configs_dict['Alice']],\n",
        "        num_rounds=1,\n",
        "    ),\n",
        "    scene_lib.SceneSpec(\n",
        "        scene_type=scene_specs['pub'],\n",
        "        start_time=PUB_TIME,\n",
        "        participant_configs=player_configs,\n",
        "        num_rounds=2,\n",
        "    ),\n",
        "    scene_lib.SceneSpec(\n",
        "        scene_type=scene_specs['decision'],\n",
        "        start_time=PUB_TIME + hour_increment,\n",
        "        participant_configs=player_configs,\n",
        "        num_rounds=1,\n",
        "    ),\n",
        "    scene_lib.SceneSpec(\n",
        "        scene_type=scene_specs['pub'],\n",
        "        start_time=PUB_TIME + (2 * hour_increment),\n",
        "        participant_configs=player_configs,\n",
        "        num_rounds=1,\n",
        "    ),\n",
        "    scene_lib.SceneSpec(\n",
        "        scene_type=scene_specs['decision'],\n",
        "        start_time=PUB_TIME + (3 * hour_increment),\n",
        "        participant_configs=player_configs,\n",
        "        num_rounds=1,\n",
        "    ),\n",
        "    scene_lib.SceneSpec(\n",
        "        scene_type=scene_specs['pub'],\n",
        "        start_time=PUB_TIME + (4 * hour_increment),\n",
        "        participant_configs=player_configs,\n",
        "        num_rounds=2,\n",
        "    ),\n",
        "]"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "metadata": {
        "id": "d2u0bQ1MSCGd"
      },
      "cell_type": "markdown",
      "source": [
        "## The RUN"
      ]
    },
    {
      "metadata": {
        "id": "hdTRDaxEZZnN"
      },
      "cell_type": "code",
      "source": [
        "clock.set(START_TIME)"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "metadata": {
        "id": "9IggLF1aH_hF"
      },
      "cell_type": "code",
      "source": [
        "for premise in scenario_premise:\n",
        "  game_master_memory.add(premise)\n",
        "  for player in players:\n",
        "    player.observe(premise)\n"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "metadata": {
        "id": "2Bt87stq76gF"
      },
      "cell_type": "code",
      "source": [
        "runner.run_scenes(\n",
        "    environment=env,\n",
        "    scenes=scenes,\n",
        "    players=players,\n",
        "    clock=clock,\n",
        ")"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "metadata": {
        "id": "DnwvpvQ4bnFs"
      },
      "cell_type": "markdown",
      "source": [
        "## Summary and analysis of the episode"
      ]
    },
    {
      "metadata": {
        "id": "5U5FDXvs4HSr"
      },
      "cell_type": "code",
      "source": [
        "# @title Metrics plotting\n",
        "\n",
        "group_by = collections.defaultdict(lambda: 'player')\n",
        "group_by['opinion_of_others'] = 'of_player'\n",
        "\n",
        "available_channels = list(measurements.available_channels())\n",
        "\n",
        "fig, ax = plt.subplots(1, len(available_channels), figsize=(6, 2))\n",
        "tb = [channel for channel in available_channels]\n",
        "for idx, channel in enumerate(available_channels):\n",
        "  plotting.plot_line_measurement_channel(measurements, channel,\n",
        "                                         group_by=group_by[channel],\n",
        "                                         xaxis='time_str',\n",
        "                                         ax=ax[idx])\n",
        "  ax[idx].set_title(channel)\n",
        "\n",
        "fig.set_constrained_layout(constrained=True)"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "metadata": {
        "id": "MuaVoPgtxzPC"
      },
      "cell_type": "markdown",
      "source": [
        "#Build and display HTML log of the experiment"
      ]
    },
    {
      "metadata": {
        "id": "j71OiuPot5UV"
      },
      "cell_type": "markdown",
      "source": [
        "## Prepare to visualize text results with HTML"
      ]
    },
    {
      "metadata": {
        "id": "O4jp0xGXvOAJ"
      },
      "cell_type": "code",
      "source": [
        "# @title Summarize the entire story.\n",
        "all_gm_memories = env._memory.retrieve_recent(k=10000, add_time=True)\n",
        "\n",
        "detailed_story = '\\n'.join(all_gm_memories)\n",
        "print('len(detailed_story): ', len(detailed_story))\n",
        "\n",
        "episode_summary = model.sample_text(\n",
        "    f'Sequence of events:\\n{detailed_story}'+\n",
        "    '\\nNarratively summarize the above temporally ordered ' +\n",
        "    'sequence of events. Write it as a news report. Summary:\\n',\n",
        "     max_tokens=3500, terminators=())\n",
        "print(episode_summary)"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "metadata": {
        "id": "ALG987t-6j-V"
      },
      "cell_type": "code",
      "source": [
        "# @title Summarise the perspective of each player\n",
        "player_logs = []\n",
        "player_log_names = []\n",
        "for player in players:\n",
        "  name = player.name\n",
        "  detailed_story = '\\n'.join(memories[player.name].retrieve_recent(\n",
        "      k=1000, add_time=True))\n",
        "  summary = ''\n",
        "  summary = model.sample_text(\n",
        "      f'Sequence of events that happened to {name}:\\n{detailed_story}'\n",
        "      '\\nWrite a short story that summarises these events.\\n',\n",
        "       max_tokens=3500, terminators=())\n",
        "\n",
        "  all_player_mem = memories[player.name].retrieve_recent(k=1000, add_time=True)\n",
        "  all_player_mem = ['Summary:', summary, 'Memories:'] + all_player_mem\n",
        "  player_html = html_lib.PythonObjectToHTMLConverter(all_player_mem).convert()\n",
        "  player_logs.append(player_html)\n",
        "  player_log_names.append(f'{name}')\n"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "metadata": {
        "id": "JyEoGgI05xI0"
      },
      "cell_type": "code",
      "source": [
        "history_sources = [env, decision_env, vodka_choice]\n",
        "histories_html = [\n",
        "    html_lib.PythonObjectToHTMLConverter(history.get_history()).convert()\n",
        "    for history in history_sources]\n",
        "histories_names = [history.name for history in history_sources]"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "metadata": {
        "id": "XNJuo4Dwt5Ui"
      },
      "cell_type": "code",
      "source": [
        "gm_mem_html = html_lib.PythonObjectToHTMLConverter(all_gm_memories).convert()\n",
        "\n",
        "tabbed_html = html_lib.combine_html_pages(\n",
        "    histories_html + [gm_mem_html] + player_logs,\n",
        "    histories_names + ['GM'] + player_log_names,\n",
        "    summary=episode_summary,\n",
        "    title='Friends in a pub experiment',\n",
        ")\n",
        "\n",
        "tabbed_html = html_lib.finalise_html(tabbed_html)"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "metadata": {
        "id": "pumxvmrzANOq"
      },
      "cell_type": "code",
      "source": [
        "# @title Display the HTML log visualization\n",
        "display.HTML(tabbed_html)"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "metadata": {
        "id": "HX-M9Im_dneG"
      },
      "cell_type": "markdown",
      "source": [
        "#Interact with a specific player"
      ]
    },
    {
      "metadata": {
        "id": "ESJ1l7_Kt5Uj"
      },
      "cell_type": "code",
      "source": [
        "sim_to_interact = 'Alice'  # @param ['Alice', 'Bob', 'Dorothy', 'Ellen'] {type:\"string\"}\n",
        "user_identity = 'a close friend'  # @param {type:\"string\"}\n",
        "interaction_premise = f'{sim_to_interact} is talking to {user_identity}\\n'  # @param {type:\"string\"}\n",
        "\n",
        "player_names = [player.name for player in players]\n",
        "player_by_name = {player.name: player for player in players}\n",
        "selected_player = player_by_name[sim_to_interact]\n",
        "interrogation = interaction_premise"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "metadata": {
        "id": "5Q1cYflLt5Uj"
      },
      "cell_type": "code",
      "source": [
        "utterence_from_user = 'Did Bob accept your apology?'  # @param {type:\"string\"}\n",
        "\n",
        "interrogation += f'{user_identity}: {utterence_from_user}'\n",
        "player_says = selected_player.say(interrogation)\n",
        "interrogation += f'\\n{sim_to_interact}: {player_says}\\n'\n",
        "print(interrogation)"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "metadata": {
        "id": "qWDqd4ByzSsT"
      },
      "cell_type": "markdown",
      "source": [
        "```\n",
        "Copyright 2023 DeepMind Technologies Limited.\n",
        "\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "you may not use this file except in compliance with the License.\n",
        "You may obtain a copy of the License at\n",
        "\n",
        "    https://www.apache.org/licenses/LICENSE-2.0\n",
        "\n",
        "Unless required by applicable law or agreed to in writing, software\n",
        "distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "See the License for the specific language governing permissions and\n",
        "limitations under the License.\n",
        "```"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "last_runtime": {
        "build_target": "",
        "kind": "private"
      },
      "private_outputs": true,
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
